-module(average).
-export([addFragmentsToAvgStruct/3, new/0, start/3, getSum/1, getCount/1, getTimeStamp/1, getFragmentIds/1, merge/2, incrTimeStamp/2]).
-import(erlang, [display/1]).

-export([test/0]).

start(Id, Neighbors, AvgStruct) ->
  random:seed(now()),
  % Number of iterations to converge
  TotalIter = ceiling(math:log(length(Neighbors)))+78,
  io:format("Total iterations ~p~n", [TotalIter]),
  NewAvgStruct = gAvg_helper(Id, Neighbors, AvgStruct, 0, TotalIter),
  io:format("~p's final avg: ~p~n", [Id, getAverage(NewAvgStruct)]),
  ok.

gAvg_helper(_Id, _Neighbors, AvgStruct, Created, TotalIter) when TotalIter+1 == Created ->
	AvgStruct;
gAvg_helper(_Id, Neighbors, AvgStruct, Created, TotalIter) ->
   case gossipprob:getGossipNode(Neighbors) of
      {ok, Node} ->
            Node ! {avg, AvgStruct, self()};
        {error, _Msg} ->
        ok
    end,
    gAvg_helper(_Id, Neighbors, processQueue(AvgStruct), Created+1, TotalIter).

processQueue(AvgStruct) ->
	receive
		ReceivedAvgStruct ->  
                     processQueue(ReceivedAvgStruct)  
	after
		12000 -> AvgStruct
	end.

	
test() ->
  D1 = dict:append_list(1, [1.0,2.0,3.0], dict:new()),
  D2 = dict:append_list(2, [4.0,5.0,6.0], D1),
  
  NewStruct = addFragmentsToAvgStruct(dict:fetch_keys(D2), D2, new()),
  display(NewStruct),
  display(getAverage(NewStruct)),
  
  NewerStruct = addFragmentsToAvgStruct(dict:fetch_keys(D2), D2, NewStruct),
  
  display(NewerStruct),
  display(getAverage(NewerStruct)),
  display("Old test"),
  
  AvgStruct = new(),
  display(AvgStruct),
  Tf = hasFragmentId(5, AvgStruct),
  display(Tf),
  AS4 = addFragmentId(5, 6, 3, AvgStruct),
  display(AS4),
  Tf2 = hasFragmentId(5, AS4),
  display(Tf2),
  AS5 = addFragmentId(4, 2, 2, AS4),
  display(AS5),
  display(hasFragmentId(4, AS5)),
  display(getSum(AS5)),
  display(getCount(AS5)),
  display(getAverage(AS5)),
  
  display("Time To Merge!"),
  Merged = merge(AS5, NewerStruct),
  display(Merged),
  display(getAverage(Merged)),
  ok.

addFragmentsToAvgStruct([], _Dict, AvgStruct) ->
  AvgStruct;
addFragmentsToAvgStruct(Keys, Dict, AvgStruct) ->
  FragId = hd(Keys),
  case hasFragmentId(FragId, AvgStruct) of
    false ->
      FloatList = dict:fetch(FragId, Dict),
      NewAvgStruct = addFragmentId(FragId, lists:sum(FloatList), length(FloatList), AvgStruct);
    true ->
      NewAvgStruct = AvgStruct
  end,
    addFragmentsToAvgStruct(tl(Keys), Dict, NewAvgStruct).
    
%Struct is a tuple that contains Sum, Count, FragmentIds -> {Sum,Count}, TimeStamp
new() ->
  {0, 0, dict:new(), 0}.
    
getAverage(AvgStruct) ->
  getSum(AvgStruct) / getCount(AvgStruct).
  
getSum(AvgStruct) ->
  element(1, AvgStruct).
  
getCount(AvgStruct) ->
  element(2, AvgStruct).
  
getFragmentIds(AvgStruct) ->
  dict:fetch_keys(getFragmentSumCounts(AvgStruct)).
  
getFragmentSumCounts(AvgStruct) ->
  element(3, AvgStruct).
  
getTimeStamp(AvgStruct) ->
  element(4, AvgStruct).
  
incrTimeStamp(Val, AvgStruct) ->
  {getSum(AvgStruct), getCount(AvgStruct), getFragmentSumCounts(AvgStruct), getTimeStamp(AvgStruct) + Val}.
  
hasFragmentId(Id, AvgStruct) ->
  lists:member(Id, getFragmentIds(AvgStruct)).


addFragmentId(Id, Sum, Count, AvgStruct) ->
  Dict = dict:store(Id, {Sum, Count}, getFragmentSumCounts(AvgStruct)),
  {getSum(AvgStruct) + Sum,
  getCount(AvgStruct) + Count,
  Dict, 
  getTimeStamp(AvgStruct)}.
  
merge(AvgStruct1, AvgStruct2) ->
  Dict1 = getFragmentSumCounts(AvgStruct1),
  RetValTmp = merge_help(dict:fetch_keys(Dict1), Dict1, new()),
  
  Dict2 = getFragmentSumCounts(AvgStruct2),
  RetVal = merge_help(dict:fetch_keys(Dict2), Dict2, RetValTmp),
  
  case getTimeStamp(AvgStruct1) < getTimeStamp(AvgStruct2) of
    true ->
      average:incrTimeStamp(getTimeStamp(AvgStruct1), RetVal);
    false ->
      average:incrTimeStamp(getTimeStamp(AvgStruct2), RetVal)
  end.
      
    
merge_help([], _SumCountDict, AvgStruct) ->
  AvgStruct;
merge_help(Keys, SumCountDict, AvgStruct) ->
  FragId = hd(Keys),
  case hasFragmentId(FragId, AvgStruct) of
    false ->
        Sum = element(1, dict:fetch(FragId, SumCountDict)),
        Count = element(2, dict:fetch(FragId, SumCountDict)),
        RetVal = addFragmentId(FragId, Sum, Count, AvgStruct);
    true ->
        RetVal = AvgStruct
    end,
    merge_help(tl(Keys), SumCountDict, RetVal).
    
    
ceiling(X) ->
    T = erlang:trunc(X),
    case (X - T) of
        Neg when Neg < 0 -> T;
        Pos when Pos > 0 -> T + 1;
        _ -> T
    end.